import json
from pathlib import Path
from agents.ddpg import DDPG


def generate_config():
    config = {
        'use_GPU': True,
        'max_steps': 3e6,
        'eval_freq': 5000,
        'batch_size': 256,
        'buffer_size': 2 ** 18,
        'start_train_steps': 25e3,
        'actor_lr': 3e-4,
        'critic_lr': 3e-4,
        'gamma': 0.99,
        'tau': 0.005,
        'expl_noise_std': 0.1,  # Std of Gaussian exploration noise
        'use_per': False,
        'alpha_per': 0.6,
        'beta_per': 0.4,
        'beta_increment': 0.001,
    }
    return config



if __name__ == "__main__":
    env_list = ['Pendulum-v1', 'BipedalWalker-v3', 'Hopper-v5', 'HalfCheetah-v5', 'Walker2d-v5', 'Swimmer-v5', 'Ant-v5',
                'Humanoid-v5']
    env = env_list[7]
    seed_list = [0, 10]
    config = generate_config()
    results = []
    for seed in seed_list:
        config['env_name'] = env
        config['seed'] = seed
        print(f"ENV:{env}  AGENT:DDPG  SEED:{seed}")
        agent = DDPG(config)
        evaluations = agent.run()
        results.append([evaluations, agent.get_score_required_to_win(), seed])
        print('-' * 100)
        print('-' * 100)
    data = {'DDPG': results}
    path = Path(f'../results/data/{env}_DDPG.json')
    json_str = json.dumps(data, indent=4)
    path.write_text(json_str)
